#pragma once

#include <linux/kernel.h>
#include <linux/cpumask_types.h>
#include <linux/bitmap.h>

/*
 * We have several different "preferred sizes" for the cpumask
 * operations, depending on operation.
 *
 * For example, the bitmap scanning and operating operations have
 * optimized routines that work for the single-word case, but only when
 * the size is constant. So if NR_CPUS fits in one single word, we are
 * better off using that small constant, in order to trigger the
 * optimized bit finding. That is 'small_cpumask_size'.
 *
 * The clearing and copying operations will similarly perform better
 * with a constant size, but we limit that size arbitrarily to four
 * words. We call this 'large_cpumask_size'.
 *
 * Finally, some operations just want the exact limit, either because
 * they set bits or just don't have any faster fixed-sized versions. We
 * call this just 'nr_cpumask_bits'.
 *
 * Note that these optional constants are always guaranteed to be at
 * least as big as 'nr_cpu_ids' itself is, and all our cpumask
 * allocations are at least that size (see cpumask_size()). The
 * optimization comes from being able to potentially use a compile-time
 * constant instead of a run-time generated exact number of CPUs.
 */
#if NR_CPUS <= BITS_PER_LONG
#define small_cpumask_bits ((unsigned int)NR_CPUS)
#define large_cpumask_bits ((unsigned int)NR_CPUS)
#elif NR_CPUS <= 4 * BITS_PER_LONG
#define small_cpumask_bits nr_cpu_ids
#define large_cpumask_bits ((unsigned int)NR_CPUS)
#else
#define small_cpumask_bits nr_cpu_ids
#define large_cpumask_bits nr_cpu_ids
#endif

/**
 * for_each_cpu - iterate over every cpu in a mask
 * @cpu: the (optionally unsigned) integer iterator
 * @mask: the cpumask pointer
 *
 * After the loop, cpu is >= nr_cpu_ids.
 */
#define for_each_cpu(cpu, mask) \
    for_each_set_bit(cpu, cpumask_bits(mask), small_cpumask_bits)

#if NR_CPUS > 1
#include <linux/smp/cpumask_mp.h>
#else
#include <linux/smp/cpumask_up.h>
#endif

void set_cpu_online(unsigned int cpu, bool possible);
void set_cpu_possible(unsigned int cpu, bool possible);
void set_cpu_present(unsigned int cpu, bool possible);

struct cpumask *cpu_present_mask(void);

/**
 * cpumask_test_cpu - test for a cpu in a cpumask
 * @cpu: cpu number (< nr_cpu_ids)
 * @cpumask: the cpumask pointer
 *
 * Return: true if @cpu is set in @cpumask, else returns false
 */
static inline bool cpumask_test_cpu(int cpu, const struct cpumask *cpumask)
{
    return test_bit(cpu, cpumask_bits((cpumask)));
}

#define cpu_is_offline(cpu) !cpu_online(cpu)

const struct cpumask *cpumask_of(unsigned int cpu);

/**
 * cpumask_copy - *dstp = *srcp
 * @dstp: the result
 * @srcp: the input cpumask
 */
static inline void cpumask_copy(struct cpumask *dstp, const struct cpumask *srcp)
{
    bitmap_copy(cpumask_bits(dstp), cpumask_bits(srcp), large_cpumask_bits);
}

/**
 * cpumask_weight - Count of bits in *srcp
 * @srcp: the cpumask to count bits (< nr_cpu_ids) in.
 *
 * Return: count of bits set in *srcp
 */
static inline unsigned int cpumask_weight(const struct cpumask *srcp)
{
    return bitmap_weight(cpumask_bits(srcp), small_cpumask_bits);
}

#if NR_CPUS <= BITS_PER_LONG
#define CPU_MASK_ALL                                                      \
    (cpumask_t)                                                           \
    {                                                                     \
        {                                                                 \
            [BITS_TO_LONGS(NR_CPUS) - 1] = BITMAP_LAST_WORD_MASK(NR_CPUS) \
        }                                                                 \
    }
#else
#define CPU_MASK_ALL                                                                                      \
    (cpumask_t)                                                                                           \
    {                                                                                                     \
        {                                                                                                 \
            [0 ... BITS_TO_LONGS(NR_CPUS) - 2] = ~0UL,                                                    \
                                            [BITS_TO_LONGS(NR_CPUS) - 1] = BITMAP_LAST_WORD_MASK(NR_CPUS) \
        }                                                                                                 \
    }
#endif /* NR_CPUS > BITS_PER_LONG */

/**
 * cpumask_first - get the first cpu in a cpumask
 * @srcp: the cpumask pointer
 *
 * Return: >= nr_cpu_ids if no cpus set.
 */
static inline unsigned int cpumask_first(const struct cpumask *srcp)
{
    return find_first_bit(cpumask_bits(srcp), small_cpumask_bits);
}

/**
 * cpumask_any - pick an arbitrary cpu from *srcp
 * @srcp: the input cpumask
 *
 * Return: >= nr_cpu_ids if no cpus set.
 */
#define cpumask_any(srcp) cpumask_first(srcp)

/**
 * to_cpumask - convert a NR_CPUS bitmap to a struct cpumask *
 * @bitmap: the bitmap
 *
 * There are a few places where cpumask_var_t isn't appropriate and
 * static cpumasks must be used (eg. very early boot), yet we don't
 * expose the definition of 'struct cpumask'.
 *
 * This does the conversion, and can be used as a constant initializer.
 */
#define to_cpumask(bitmap)           \
    ((struct cpumask *)(1 ? (bitmap) \
                          : (void *)sizeof(__check_is_bitmap(bitmap))))

static inline int __check_is_bitmap(const unsigned long *bitmap)
{
    return 1;
}

/**
 * cpumask_set_cpu - set a cpu in a cpumask
 * @cpu: cpu number (< nr_cpu_ids)
 * @dstp: the cpumask pointer
 */
static inline void cpumask_set_cpu(unsigned int cpu, struct cpumask *dstp)
{
    set_bit(cpu, cpumask_bits(dstp));
}
